package xyz.scootaloo.code.adt;

import org.jetbrains.annotations.NotNull;

/**
 * 二分搜索树
 *
 * @author flutterdash@qq.com
 * @since 2021/11/10 13:46
 */
public class BinarySearchTree<E extends Comparable<E>>
        extends AbstractBinarySearchTree<E> implements KCollection<E> {

    private int size;
    private BinaryNode<E> root;

    public BinarySearchTree() {
        clear();
    }

    @Override
    public int size() {
        return this.size;
    }

    @Override
    public void clear() {
        size = 0;
        root = null;
    }

    @Override
    public boolean add(@NotNull E element) {
        if (contains(element))
            return false;

        if (root == null) {
            root = createNode(element);
            size++;
            return true;
        }

        binaryInsert(root, element);
        size++;
        return true;
    }

    @Override
    public boolean remove(@NotNull E element) {
        if (contains(element))
            return false;
        root = binaryRemove(root, element);
        size--;
        return true;
    }

    @Override
    public boolean contains(@NotNull E element) {
        return binarySearch(root, element) != null;
    }

    @NotNull
    @Override
    public KIterator<E> iterator() {
        return treeIterator(root);
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    /*
     ================================
     =                              =
     =   二叉搜索树算法实现细节部分
     =                              =
     ================================
     */

    /**
     * 查找目标元素是否存在于当前节点或其子树中
     * <p/>
     * 二分查找根据调用元素的compareTo方法实现, 此方法返回一个int值;
     * 1. diff == 0: 当前节点的元素值与目标元素一致;
     * 2. diff < 0: 当前节点的元素值比目标元素值小, 因此目标节点应该处于当前节点的右子树中;
     * 3. diff > 0: 当前节点的元素值比目标元素值大, 因此目标节点应该处于当前节点的左子树中;
     *
     * @param current 当前节点
     * @param target  目标元素
     * @return 节点; 空(假如集合中不存在此元素)
     */
    private BinaryNode<E> binarySearch(BinaryNode<E> current, E target) {
        if (current == null)
            return null;

        int diff = current.value.compareTo(target);
        if (diff == 0) { // 1
            return current;
        } else if (diff < 0) { // 2
            return binarySearch(current.right, target);
        } else { // 3
            return binarySearch(current.left, target);
        }
    }

    /**
     * 将一个目标元素插入到当前节点的子树中
     * <p>
     * ! 插入前请确保root不为空;
     * <p>
     * 插入条件:
     * 1. 目标元素值与当前节点的元素值一致, 无需插入;
     * 2. 目标元素值比当前节点值大, 且当前节点的右孩子节点恰好为空, 将目标元素插入到右孩子节点;
     * 3. 目标元素值比当前节点值小, 且当前节点的左孩子节点恰好为空, 将目标元素插入到左孩子节点;
     *
     * @param current 当前节点
     * @param element 目标值
     */
    private BinaryNode<E> binaryInsert(BinaryNode<E> current, E element) {
        if (current == null)
            return createNode(element);

        int diff = current.value.compareTo(element);
        if (diff < 0) { // 2
            current.right = binaryInsert(current.right, element);
        } else { // 1, 3
            current.left = binaryInsert(current.left, element);
        }

        return current;
    }

    /**
     * 删除父节点的一个子孙节点
     * <p>
     * 1. 要被删除的节点没有左孩子和右孩子;
     * 2. 要被删除的节点只有左孩子;
     * 3. 要被删除的节点只有右孩子;
     * 4. 要被删除的节点同时有左孩子和右孩子;
     *
     * @param current 当前节点, 目标元素所在的节点在当前节点或其子孙节点中
     * @param element 目标元素
     * @return 假如子孙节点确实是父节点的子孙节点, 并且成功切断这个节点与整个集合的连接, 则删除成功
     */
    private BinaryNode<E> binaryRemove(BinaryNode<E> current, E element) {
        if (current == null)
            return null;

        int diff = current.value.compareTo(element);
        /* diff > 0 或 diff < 0, 要被删除的节点在当前节点的子孙节点中 */
        if (diff < 0) {
            current.left = binaryRemove(current.left, element);
        } else if (diff > 0) {
            current.right = binaryRemove(current.right, element);
        }
        /* 其余情况代表当前节点就是要被删除的节点, diff == 0 */
        else if (current.left != null && current.right != null) { // 4: 左右孩子节点都不为空
            current.value = findMin(current).value;
            current.right = binaryRemove(current.right, current.value);
        } else { // 1, 2, 3
            current = (current.left != null) ? current.left : current.right;
        }

        return current;
    }

    private BinaryNode<E> createNode(E element) {
        return new BinaryNode<>(element);
    }
}
